# src/augmentations.py

"""
时间序列数据增强模块
为HRV时间序列提供多种基于信号的增强方法
"""

import numpy as np
from scipy.interpolate import CubicSpline
import random


class TimeSeriesAugmentation:
    """
    时间序列数据增强类
    提供多种增强策略以提升模型的泛化能力
    """
    
    def __init__(self, augmentation_config=None):
        """
        Args:
            augmentation_config: 增强配置字典，包含各种增强方法的参数
        """
        if augmentation_config is None:
            # 默认配置
            self.config = {
                'jitter_sigma': 0.03,
                'scale_sigma': 0.1,
                'magnitude_warp_knots': 4,
                'magnitude_warp_sigma': 0.2,
                'time_warp_knots': 4,
                'time_warp_sigma': 0.2,
                'apply_probability': 0.5
            }
        else:
            self.config = augmentation_config
    
    def jitter(self, x, sigma=None):
        """
        添加高斯噪声（抖动）
        
        Args:
            x: 输入序列 (sequence_length, num_features)
            sigma: 噪声标准差（默认使用配置值）
        
        Returns:
            增强后的序列
        """
        if sigma is None:
            sigma = self.config['jitter_sigma']
        
        noise = np.random.normal(loc=0., scale=sigma, size=x.shape)
        return x + noise
    
    def scaling(self, x, sigma=None):
        """
        随机缩放序列的幅度
        
        Args:
            x: 输入序列 (sequence_length, num_features)
            sigma: 缩放因子的标准差（默认使用配置值）
        
        Returns:
            增强后的序列
        """
        if sigma is None:
            sigma = self.config['scale_sigma']
        
        # 为每个特征生成独立的缩放因子
        scaling_factors = np.random.normal(loc=1.0, scale=sigma, size=(1, x.shape[1]))
        return x * scaling_factors
    
    def magnitude_warping(self, x, n_knots=None, sigma=None):
        """
        幅度扭曲：通过平滑的三次样条曲线扭曲序列的幅度
        
        Args:
            x: 输入序列 (sequence_length, num_features)
            n_knots: 样条曲线的节点数（默认使用配置值）
            sigma: 扭曲强度（默认使用配置值）
        
        Returns:
            增强后的序列
        """
        if n_knots is None:
            n_knots = self.config['magnitude_warp_knots']
        if sigma is None:
            sigma = self.config['magnitude_warp_sigma']
        
        seq_len, num_features = x.shape
        x_warped = np.zeros_like(x)
        
        # 为每个特征独立应用幅度扭曲
        for i in range(num_features):
            # 生成随机节点
            knot_positions = np.linspace(0, seq_len - 1, n_knots + 2)
            knot_values = np.random.normal(loc=1.0, scale=sigma, size=n_knots + 2)
            
            # 固定端点以保持序列的整体范围
            knot_values[0] = 1.0
            knot_values[-1] = 1.0
            
            # 创建平滑的三次样条曲线
            cs = CubicSpline(knot_positions, knot_values)
            
            # 生成整个序列的扭曲因子
            warp_factors = cs(np.arange(seq_len))
            
            # 应用扭曲
            x_warped[:, i] = x[:, i] * warp_factors
        
        return x_warped
    
    def time_warping(self, x, n_knots=None, sigma=None):
        """
        时间扭曲：通过非线性时间扭曲来拉伸或压缩序列
        
        Args:
            x: 输入序列 (sequence_length, num_features)
            n_knots: 样条曲线的节点数（默认使用配置值）
            sigma: 扭曲强度（默认使用配置值）
        
        Returns:
            增强后的序列
        """
        if n_knots is None:
            n_knots = self.config['time_warp_knots']
        if sigma is None:
            sigma = self.config['time_warp_sigma']
        
        seq_len, num_features = x.shape
        
        # 生成随机时间扭曲路径
        knot_positions = np.linspace(0, seq_len - 1, n_knots + 2)
        
        # 生成随机扭曲值
        random_warps = np.random.normal(loc=0.0, scale=sigma * seq_len, size=n_knots + 2)
        
        # 固定起始和结束点
        random_warps[0] = 0
        random_warps[-1] = 0
        
        # 生成目标时间点（保证单调递增）
        warped_knots = knot_positions + random_warps
        warped_knots = np.sort(warped_knots)
        
        # 确保在有效范围内
        warped_knots[0] = 0
        warped_knots[-1] = seq_len - 1
        
        # 创建时间扭曲映射
        time_warp = CubicSpline(knot_positions, warped_knots)
        
        # 生成新的时间索引
        new_indices = time_warp(np.arange(seq_len))
        new_indices = np.clip(new_indices, 0, seq_len - 1)
        
        # 对每个特征进行插值
        x_warped = np.zeros_like(x)
        original_indices = np.arange(seq_len)
        
        for i in range(num_features):
            # 使用线性插值重新采样
            x_warped[:, i] = np.interp(new_indices, original_indices, x[:, i])
        
        return x_warped
    
    def apply_random_augmentation(self, x, p=None):
        """
        随机选择一种或多种增强方法并应用
        
        Args:
            x: 输入序列 (sequence_length, num_features)
            p: 应用每种增强的概率（默认使用配置值）
        
        Returns:
            增强后的序列
        """
        if p is None:
            p = self.config['apply_probability']
        
        augmentations = [
            ('jitter', self.jitter),
            ('scaling', self.scaling),
            ('magnitude_warping', self.magnitude_warping),
            ('time_warping', self.time_warping)
        ]
        
        # 创建副本以避免修改原始数据
        x_aug = x.copy()
        
        # 随机应用增强
        applied_augs = []
        for name, aug_func in augmentations:
            if np.random.random() < p:
                x_aug = aug_func(x_aug)
                applied_augs.append(name)
        
        # 如果没有应用任何增强，至少应用一种
        if len(applied_augs) == 0:
            chosen_aug = random.choice(augmentations)
            x_aug = chosen_aug[1](x_aug)
        
        return x_aug
    
    def apply_composed_augmentation(self, x, augmentation_list):
        """
        按顺序应用指定的增强方法
        
        Args:
            x: 输入序列 (sequence_length, num_features)
            augmentation_list: 增强方法名称列表，例如 ['jitter', 'scaling']
        
        Returns:
            增强后的序列
        """
        augmentation_map = {
            'jitter': self.jitter,
            'scaling': self.scaling,
            'magnitude_warping': self.magnitude_warping,
            'time_warping': self.time_warping
        }
        
        x_aug = x.copy()
        
        for aug_name in augmentation_list:
            if aug_name in augmentation_map:
                x_aug = augmentation_map[aug_name](x_aug)
            else:
                print(f"警告: 未知的增强方法 '{aug_name}'，已跳过")
        
        return x_aug


def get_augmentation(config=None):
    """
    工厂函数：创建并返回TimeSeriesAugmentation实例
    
    Args:
        config: 增强配置字典
    
    Returns:
        TimeSeriesAugmentation实例
    """
    return TimeSeriesAugmentation(config)


# ================== 示例和测试 ==================
if __name__ == "__main__":
    import matplotlib.pyplot as plt
    
    # 生成模拟的HRV时间序列数据
    np.random.seed(42)
    seq_len = 256
    num_features = 3
    
    # 创建一个简单的测试序列（正弦波 + 噪声）
    t = np.linspace(0, 4 * np.pi, seq_len)
    x_test = np.zeros((seq_len, num_features))
    x_test[:, 0] = np.sin(t) + 0.1 * np.random.randn(seq_len)
    x_test[:, 1] = np.cos(t) + 0.1 * np.random.randn(seq_len)
    x_test[:, 2] = np.sin(2 * t) + 0.1 * np.random.randn(seq_len)
    
    # 创建增强器
    augmentor = TimeSeriesAugmentation()
    
    # 测试各种增强方法
    fig, axes = plt.subplots(5, 1, figsize=(12, 15))
    
    # 原始信号
    axes[0].plot(x_test[:, 0], label='Feature 1')
    axes[0].plot(x_test[:, 1], label='Feature 2')
    axes[0].plot(x_test[:, 2], label='Feature 3')
    axes[0].set_title('Original Signal')
    axes[0].legend()
    axes[0].grid(True)
    
    # Jittering
    x_jitter = augmentor.jitter(x_test)
    axes[1].plot(x_jitter[:, 0], label='Feature 1')
    axes[1].plot(x_jitter[:, 1], label='Feature 2')
    axes[1].plot(x_jitter[:, 2], label='Feature 3')
    axes[1].set_title('Jittering (Gaussian Noise)')
    axes[1].legend()
    axes[1].grid(True)
    
    # Scaling
    x_scale = augmentor.scaling(x_test)
    axes[2].plot(x_scale[:, 0], label='Feature 1')
    axes[2].plot(x_scale[:, 1], label='Feature 2')
    axes[2].plot(x_scale[:, 2], label='Feature 3')
    axes[2].set_title('Scaling')
    axes[2].legend()
    axes[2].grid(True)
    
    # Magnitude Warping
    x_mag_warp = augmentor.magnitude_warping(x_test)
    axes[3].plot(x_mag_warp[:, 0], label='Feature 1')
    axes[3].plot(x_mag_warp[:, 1], label='Feature 2')
    axes[3].plot(x_mag_warp[:, 2], label='Feature 3')
    axes[3].set_title('Magnitude Warping')
    axes[3].legend()
    axes[3].grid(True)
    
    # Time Warping
    x_time_warp = augmentor.time_warping(x_test)
    axes[4].plot(x_time_warp[:, 0], label='Feature 1')
    axes[4].plot(x_time_warp[:, 1], label='Feature 2')
    axes[4].plot(x_time_warp[:, 2], label='Feature 3')
    axes[4].set_title('Time Warping')
    axes[4].legend()
    axes[4].grid(True)
    
    plt.tight_layout()
    plt.savefig('augmentation_examples.png', dpi=150, bbox_inches='tight')
    print("增强示例图已保存到: augmentation_examples.png")
    
    print("\n✅ 所有增强方法测试通过！")


